#!/bin/bash

set -e
pushd $(cd $(dirname ${0})/..; pwd) > /dev/null

case "${OSTYPE}" in
    msys*) python="winpty python";;
    *) python="python";;
esac

venv() {
    env_dir=./third_party/env/onnx
    [ -d "${env_dir}" ] || ${python} -m venv ${env_dir}
    case "${OSTYPE}" in
        msys*) source ${env_dir}/Scripts/activate;;
        *) source ${env_dir}/bin/activate;;
    esac
    ${python} -m pip install --quiet --upgrade pip setuptools wheel
}

clean() {
    echo "sklearn clean"
    rm -rf "./third_party/env/onnx"
    rm -rf "./third_party/source/onnx"
}

sync() {
    echo "onnx sync"
    [ -d "./third_party/source/onnx" ] || git clone --quiet --recursive https://github.com/onnx/onnx.git "./third_party/source/onnx"
    pushd "./third_party/source/onnx" > /dev/null
    git pull --quiet --prune
    git submodule sync --quiet
    git submodule update --quiet --init --recursive
    popd > /dev/null
}

install() {
    echo "onnx install"
    case "${OSTYPE}" in
        linux*)
            [ -x "$(command -v cmake)" ] || sudo apt install -y cmake
            [ -x "$(command -v protoc)" ] || sudo apt install -y protobuf-compiler libprotoc-dev
            protobuf=protobuf
            ;;
        darwin*)
            [ -x "$(command -v cmake)" ] || brew install cmake > /dev/null
            [ -x "$(command -v protoc)" ] || brew install protobuf > /dev/null
            protobuf=protobuf
            if [ "$(uname -sm)" = "Darwin arm64" ]; then
                PYTHON_ARCHITECTURE=$(${python} -c "import platform; print(platform.uname().machine);")
                export CMAKE_ARGS="-DCMAKE_OSX_ARCHITECTURES=${PYTHON_ARCHITECTURE}"
                export OPENBLAS=$(brew --prefix openblas)
            fi
            ;;
        msys*)
            [ ! -z "$(choco list --local-only --exacty --limit-output visualstudio2017-workload-vctools)" ] || $(choco install -yes visualstudio2017-workload-vctools) > /dev/null
            protoc_version=3.9.x
            protoc_dir="$(pwd)/third_party/bin/protobuf/v${protoc_version}"
            programfiles_x86_dir=$(env | grep "^ProgramFiles(x86)=" | cut -d '=' -f 2)
            cmake_dir="${programfiles_x86_dir}\Microsoft Visual Studio\2017\BuildTools\Common7\IDE\CommonExtensions\Microsoft\CMake\CMake\bin"
            msbuild_dir="${programfiles_x86_dir}\Microsoft Visual Studio\2017\BuildTools\MSBuild\15.0\Bin"
            if [ ! -f "${protoc_dir}/bin/protoc.exe" ]; then
                rm -rf ${protoc_dir}
                git clone --quiet --branch ${protoc_version} https://github.com/protocolbuffers/protobuf.git ${protoc_dir}/src
                pushd "${protoc_dir}/src/cmake" > /dev/null
                "${cmake_dir}\cmake.exe" -G "Visual Studio 15 2017 Win64" -Dprotobuf_MSVC_STATIC_RUNTIME=OFF -Dprotobuf_BUILD_TESTS=OFF -Dprotobuf_BUILD_EXAMPLES=OFF -DCMAKE_INSTALL_PREFIX="..\.." > /dev/null
                "${msbuild_dir}\MSBuild.exe" protobuf.sln //m //p:Configuration=Release > /dev/null
                "${msbuild_dir}\MSBuild.exe" INSTALL.vcxproj //p:Configuration=Release > /dev/null
                popd > /dev/null
            fi
            export PATH="${protoc_dir}\bin":"$(cygpath -u "${cmake_dir}")":${PATH}
            export USE_MSVC_STATIC_RUNTIME=0
            protobuf="protobuf==3.9.2"
            ;;
    esac
    venv
    ${python} -m pip install --quiet --upgrade ${protobuf}
    ${python} -m pip install --quiet --upgrade numpy --no-use-pep517
    export ONNX_ML=1
    export ONNX_NAMESPACE=onnx
    ${python} -m pip install --quiet "./third_party/source/onnx"
    deactivate
}

schema() {
    echo "onnx schema"
    [[ $(grep -U $'\x0D' ./source/onnx-proto.js) ]] && crlf=1
    node ./tools/protoc.js --text --root onnx --out ./source/onnx-proto.js --path ./third_party/source/onnx onnx/onnx-ml.proto onnx/onnx-operators-ml.proto
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/onnx-proto.js ./source/onnx-proto.js
    fi
}

metadata() {
    echo "onnx metadata"
    [[ $(grep -U $'\x0D' ./source/onnx-metadata.json) ]] && crlf=1
    venv
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py metadata
    deactivate
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/onnx-metadata.json ./source/onnx-metadata.json
    fi
}

infer() {
    echo "onnx infer"
    venv
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py infer ${1}
    deactivate
}

optimize() {
    echo "onnx optimize"
    venv
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py optimize ${1}
    deactivate
}

while [ "$#" != 0 ]; do
    command="$1" && shift
    case "${command}" in
        "clean") clean;;
        "sync") sync;;
        "install") install;;
        "schema") schema;;
        "metadata") metadata;;
        "infer") infer ${1} && shift;;
        "optimize") optimize ${1} && shift;;
    esac
done
